import platgo as pg
import numpy as np
from numpy import sum, sin, pi, square, sqrt, mean


class UserProblem(pg.Problem):

    def __init__(self, num_of_obj: str = 'multi', encoding: str = 'real', other: list = None, D: int = 10,
                 ub: float = np.inf, lb: float = -np.inf, cal_obj: list = None, cal_con: list = None):
        """ initFcn;                % function for initializing a population
        decFcn;                 % function for repairing invalid solution
        objFcn = @(x,d)sum(x);	% objective functions
        conFcn = @(x,d)0;      	% constraint functions """

        self.type[num_of_obj] = True
        self.type[encoding] = True
        if other:
            for item in other:
                self.type[item] = True
        self.name = 'UserProblem'
        self.calobj = cal_obj
        self.calcon = cal_con
        self.D = D
        self.M = len(cal_obj)
        self.N = 100
        lb = [lb] * self.D
        ub = [ub] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def init_pop(self, N: int = None) -> pg.Population:
        if not N:
            pass
        else:
            self.N = N
        if self.type['binary']:
            decs = np.random.randint(low=0, high=2, size=(N, self.D))
        elif self.type['permutation']:
            decs = np.argsort(np.random.random(size=(self.N, self.D)), axis=1)
        else:
            decs = np.random.uniform(np.tile(self.borders[0], (self.N, 1)), np.tile(self.borders[1], (self.N, 1)))
        pop = pg.Population(decs=decs)
        return pop

    def cal_obj(self, pop) -> None:

        if self.M == 1:
            pop.objv = np.array([eval(self.calobj[0].replace('x', 'pop.decs'))]).T
        else:
            for item in self.calobj:
                objv = eval(item.replace('x', 'pop.decs'))
                if pop.objv is not None:
                    pop.objv = np.hstack((pop.objv, objv))
                else:
                    pop.objv = objv
        if self.type['constrained']:
            pop.cv = np.array([eval(self.calcon[0])]).T
        else:
            pass

    def get_optimal(self) -> np.ndarray:
        pass



